#!/usr/bin/python3
#ansible python api

import json
import shutil
#用于添加选项，比如指定远程用户 remote_user=None
from ansible.module_utils.common.collections import ImmutableDict
#读取json/yaml/ini格式文件的数据解析器
from ansible.parsing.dataloader import DataLoader
#管理主机和主机组的变量管理器
from ansible.vars.manager import VariableManager
#管理资源库的，可指定一个inventory文件
from ansible.inventory.manager import InventoryManager
#用于执行ad-hoc类，需要传入相应参数
from ansible.playbook.play import Play
#ansible底层用到的任务队列管理器
from ansible.executor.task_queue_manager import TaskQueueManager
#处理任务执行后返回的状态
from ansible.plugins.callback import CallbackBase
#上下文管理器，用来接收ImmutableDict 的示例对象
from ansible import context
#用户获取ansible产生的临时文档
import ansible.constants as C

#用于执行playbook核心类
from ansible.executor.playbook_executor import PlaybookExecutor
#主机组操作，可以给组添加变量等操作，扩展
from ansible.inventory.host import Group
#对主机进行操作，可以给主机添加变量等操作，扩展
from ansible.inventory.host import Host

class ResultCallback(CallbackBase):
    '''
    回调插件，用于对执行结果的回调
    如果要将执行的命令的所有结果放到一个对象里，应该看看如何使用json回调插件
    或者编写自定义的回调插件
    '''
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.host_ok = {}
        self.host_unreachable = {}
        self.host_failed = {}

    def v2_runner_on_ok(self, result, **kwargs):
        '''
        将结果以json的格式打印出来
        此方法可以将结果存储在实例属性中以便稍后检索 
        '''
        self.host_ok[result._host.get_name()] = result
   
    def v2_runner_on_unreachable(self, result):
        self.host_unreachable[result._host.get_name()] = result

    def v2_runner_on_failed(self, result, **kwargs):
        self.host_failed[result._host.get_name()] = result

class MyAnsible():
    def __init__(self,
        connection='local', #连接方式local本地方式，smart ssh方式
        remote_user=None, #ssh用户
        remote_password=None, #ssh用户密码。应该是一个字典，key必须是conn_pass
        private_key_file=None, #指定自定义私钥
        sudo=None, sudo_user=None, ask_sudo_pass=None,
        module_path=None, #模块路径，可以指定一个自定义模块的路径
        become=None, #是否提权
        become_method=None, #提权方式 默认sudo可以是su      
        become_user=None, #提权后，要成为的用户，并非登录用户
        listhosts=None, listtasks=None, listtags=None,
        verbosity=3, syntax=None, start_at_task=None,
        inventory=None, extra_vars={}):
        '''
        初始化函数，定义的默认选项值
        在初始化的时候可以传参，以便覆盖默认选项的值
        '''
        context.CLIARGS = ImmutableDict(
            connection=connection, #连接方式local本地方式，smart ssh方式
            remote_user=remote_user, #ssh用户
            private_key_file=private_key_file, #指定自定义私钥
            sudo=sudo, sudo_user=sudo_user, ask_sudo_pass=ask_sudo_pass,
            module_path=module_path, #模块路径，可以指定一个自定义模块的路径
            become=become, #是否提权
            become_method=become_method, #提权方式 默认sudo可以是su      
            become_user=become_user, #提权后，要成为的用户，并非登录用户
            listhosts=listhosts, listtasks=listtasks, listtags=listtags,
            verbosity=verbosity, syntax=syntax, start_at_task=start_at_task,
        )
        
        #三元表达式，如果没有传递inventory，就使用'localhot,'
        #指定inventory文件
        #inventory的值可以是一个资产清单文件
        #也可以是一个包含主机的元祖，一般仅仅用于测试
        #比如1.1.1.1 如果只有一个ip最后必须有英文的逗号
        #或者 1.1.1.1,2.2.2.2

        self.inventory = inventory if inventory else "localhost,"

        #实例化数据解析器
        self.loader = DataLoader()
      
        #实例化资产配置对象
        self.inv_obj = InventoryManager(loader=self.loader, sources=self.inventory)

        #设置密码
        self.passwords = remote_password
        
        #实例化回调插件对象
        self.results_callback = ResultCallback()

        #变量管理器
        self.variable_manager = VariableManager(self.loader, self.inv_obj)
       
        #自定义变量
        self.extra_vars = extra_vars
   
    def run(self, hosts='localhost', gather_facts="no", module="ping", args='', task_time=0):
        '''
        参数说明:
        task_time -- 执行异步任务时等待的秒数，这个需要大于0，等于0不支持异步（默认值）
        这个值应该等于执行任务实际耗时间尾号
        '''
        play_source = dict(
            name = "Ad-hoc",
            hosts = hosts,
            gather_facts = gather_facts,
            tasks = [
                #这里每个task就是这个列表中的一个元素，格式是嵌套的字典
                {"action": {"module": module, "args": args}, "async": task_time, "poll": 0}
            ]
        )

        play = Play().load(play_source, variable_manager=self.variable_manager, loader=self.loader)
        
        tqm = None
  
        try:
            tqm = TaskQueueManager(
                inventory=self.inv_obj,
                variable_manager=self.variable_manager,
                loader=self.loader,
                passwords=self.passwords,
                stdout_callback=self.results_callback
            )
            result = tqm.run(play)
        finally:
            if tqm is not None:
                tqm.cleanup()

      
            shutil.rmtree(C.DEFAULT_LOCAL_TMP, True)

    def playbook(self, playbooks):
        '''
        keyword arguments
        playbooks -- 需要是一个列表类型
        '''
        from ansible.executor.playbook_executor import PlaybookExecutor

        self.variable_manager._extra_vars = self.extra_vars
        playbook = PlaybookExecutor(
            playbooks=playbooks,
            inventory=self.inv_obj,
            variable_manager=self.variable_manager,
            loader=self.loader,
            passwords=self.passwords
        )
        
        #使用回调函数
        playbook._tqm._stdout_callback = self.results_callback

        result = playbook.run()

    def get_result(self):
        result_raw = {'success': {}, 'failed': {}, 'unreachable': {}} 
        
        for host, result in self.results_callback.host_ok.items():
            result_raw['success'][host] = result._result

        for host, result in self.results_callback.host_failed.items():
            result_raw['failed'][host] = result._result

        for host, result in self.results_callback.host_unreachable.items():
            result_raw['unreachable'][host] = result._result
        
        return json.dumps(result_raw, indent=4)

#执行默认的Ad-hoc
#ansible2 = MyAnsible()
#ansible2.run()
#ansible2.get_result()

#执行自定义的Ad-hoc
#使用自己的资产配置文件，并使用ssh远程连接方式
#ansible2 = MyAnsible(inventory='/etc/ansible/hosts', connection='smart')
#执行自定义任务，执行对象是zabbix组
#ansible2.run(hosts='zabbix', module="shell", args="ifconfig")
#打印结果
#ansible2.get_result()

#执行playbook
#ansible2 = MyAnsible(inventory='/etc/ansible/hosts', connection='smart', extra_vars={"dir": "/opt/scripts"})
#传入playbooks参数，需要的是一个列表数据类型，这里使用相对路径
#相对路径是相对于执行脚本的当前用户的家目录
#ansible2.playbook(playbooks=['test_ping.yml'])
#ansible2.get_result()

#执行异步任务
#ansible2 = MyAnsible(inventory='/etc/ansible/hosts', connection='smart')
#ansible2.run(module="shell", args="sleep 15;hostname -i", task_time=15)
#ansible2.get_result()

